import time
import av
import collections


import yaml
from  tools.lock.file_lock import FileLock
from tools.log import logger
from concurrent.futures import ThreadPoolExecutor
from tools.log import logger
class H264Camera:
    def __init__(self, conf: dict) -> None:
        self.url = conf["url"]
        self.max_seconds = conf["ring"]["length"]
        self.container = None
        self.video_stream = None
        self.time_base = None
        self.ring_buffer = None
        self.pts_sec = 0
        self.executor = ThreadPoolExecutor(max_workers=2)
        # 初始化连接
        self.connect_with_retry()

    def connect_with_retry(self):
        """
        无限重连 RTSP，直到成功打开流。
        """
        while True:
            try:
                logger.info(f"尝试连接 RTSP: {self.url}")
                self.container = av.open(
                    self.url,
                    options={
                        "rtsp_transport": "tcp",
                        "stimeout": str(5 * 1_000_000),  # 5 秒
                        "framedrop": "1",
                        "max_delay": "500000",  # 0.5 秒
                        "hwaccel": "rkmpp"
                    }
                )
                self.video_stream = next(s for s in self.container.streams if s.type == 'video')
                self.time_base = self.video_stream.time_base
                # 初始化环形缓冲区
                self.ring_buffer = H264RingBuffer(self.max_seconds, self.time_base)

                logger.info(f"成功连接 RTSP: {self.url}")
                break  # 成功打开流，退出重连循环
            except Exception as e:
                logger.error(f"连接失败: {e}, 1 秒后重试")
                time.sleep(1)

               
    def add_pkts(self,packet):
        if packet.dts is None:
            return False
        self.ring_buffer.append_packet(packet)
        self.pts_sec = float(packet.pts *self.time_base)
        return True
    
    def export_mp4(self,d_time,meta,pending_file,pending_lock):
        """
        output_filename must be .mp4
        """
        selected_packets = self.ring_buffer.export_pkts(self.pts_sec-d_time,self.pts_sec)
        # threading.Thread(target=_export_mp4,args=(selected_packets,meta,pending_file,pending_lock)).start()
        self.executor.submit(_export_mp4, selected_packets, meta, pending_file, pending_lock)
        return
    
def write_meta(file,meta_yaml):
    with open(file,"w",encoding='utf-8') as wd:
        yaml.safe_dump(meta_yaml,wd)    

def _export_mp4(selected_packets,meta,pending_file,pending_lock):
    """
    output_filename must be .mp4
    """
    logger.info(f'video {meta["video"]} 开始导出')
    output_container = av.open(meta["dir"]+meta["video"], mode='w')
    out_stream = output_container.add_stream('hevc')

    base_pts = selected_packets[0].pts
    base_dts = selected_packets[0].dts
    for pkt in selected_packets:

        pkt.pts = pkt.pts - base_pts
        pkt.dts = pkt.dts - base_dts
        pkt.stream = out_stream
        output_container.mux(pkt)
    output_container.close()
    logger.info(f'video {meta["video"]} 导出成功')  
    logger.info(f"写入元数据{meta['video']}...")
    write_meta(meta["dir"] +"meta.yaml",meta)  
    logger.info(f"完成元数据{meta['video']}")
    lock = FileLock(pending_lock)
    with lock:
        with open(pending_file,"a",encoding="utf-8") as file:
            file.write(meta["dir"]+"\n")
    logger.info(f"{meta['id']} 通知完毕")

    
class H264RingBuffer:
    def __init__(self, max_seconds, time_base):
        self.max_seconds = max_seconds
        self.time_base = time_base
        self.buffer = collections.deque()  # 存储 (pts_sec, packet, is_keyframe)

    def append_packet(self, packet):
        # 计算时间戳秒数
        pts_sec = float(packet.pts * self.time_base)
        # 简单时间戳容错：保证单调递增
        if self.buffer and pts_sec < self.buffer[-1][0]:
            pts_sec = self.buffer[-1][0]

        is_keyframe = packet.is_keyframe

        self.buffer.append((pts_sec, packet, is_keyframe))
        self._trim_buffer(pts_sec)

    def _trim_buffer(self, current_pts_sec):
        # 保证缓冲区只保留 max_seconds 秒内的数据
        while self.buffer and (current_pts_sec - self.buffer[0][0]) > self.max_seconds:
            self.buffer.popleft()

    def export_pkts(self, start_sec, end_sec):
        keyframe_index = None
        for i, (pts, pkt, is_key) in enumerate(self.buffer):
            if pts >= start_sec and is_key:
                keyframe_index = i
                break
        if keyframe_index is None:
            keyframe_index = 0

        selected_packets = []
        for pts, pkt, _ in list(self.buffer)[keyframe_index:]:
            if pts > end_sec:
                break
            new_pkt = av.packet.Packet(bytes(pkt))
            new_pkt.pts = pkt.pts
            new_pkt.dts = pkt.dts
            new_pkt.time_base = pkt.time_base
            selected_packets.append(new_pkt)

        if not selected_packets:
            logger.info("指定时间段内无可导出包")
            return None
        return selected_packets

